package com.sun.dream.fish.secret.mac;

import org.bouncycastle.crypto.engines.ChaChaEngine;
import org.bouncycastle.crypto.macs.Poly1305;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.Arrays;

import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;

public class ChaCha20Poly1305Example {
    public static void main(String[] args) {
        String plaintext = "Hello World";
        byte[] key = generateRandomKey(); // Generate a 256-bit random key
        byte[] nonce = generateRandomNonce(); // Generate a random 96-bit nonce

        byte[] ciphertext = encrypt(plaintext.getBytes(StandardCharsets.UTF_8), key, nonce);

        String decryptedText = decrypt(ciphertext, key, nonce);

        System.out.println("Origin Text: " + plaintext);
        System.out.println("Encrypted text: " + bytesToHex(ciphertext));
        System.out.println("Decrypted Text: " + decryptedText);
    }

    public static byte[] encrypt(byte[] plaintext, byte[] key, byte[] nonce) {
        ChaChaEngine chacha20 = new ChaChaEngine();
        chacha20.init(true, new ParametersWithIV(new KeyParameter(key), nonce));

        byte[] ciphertext = new byte[plaintext.length];
        chacha20.processBytes(plaintext, 0, plaintext.length, ciphertext, 0);

        Poly1305 poly1305 = new Poly1305();
        poly1305.init(new KeyParameter(key));

        poly1305.update(ciphertext, 0, ciphertext.length);
        byte[] mac = new byte[poly1305.getMacSize()];
        poly1305.doFinal(mac, 0);

        return Arrays.concatenate(ciphertext, mac);
    }

    public static String decrypt(byte[] ciphertext, byte[] key, byte[] nonce) {
        byte[] receivedCiphertext = Arrays.copyOfRange(ciphertext, 0, ciphertext.length - 16);
        byte[] receivedMac = Arrays.copyOfRange(ciphertext, ciphertext.length - 16, ciphertext.length);

        Poly1305 poly1305 = new Poly1305();
        poly1305.init(new KeyParameter(key));

        poly1305.update(receivedCiphertext, 0, receivedCiphertext.length);
        byte[] computedMac = new byte[poly1305.getMacSize()];
        poly1305.doFinal(computedMac, 0);

        if (!Arrays.areEqual(receivedMac, computedMac)) {
            throw new IllegalArgumentException("Invalid MAC");
        }

        ChaChaEngine chacha20 = new ChaChaEngine();
        chacha20.init(false, new ParametersWithIV(new KeyParameter(key), nonce));

        byte[] decryptedText = new byte[receivedCiphertext.length];
        chacha20.processBytes(receivedCiphertext, 0, receivedCiphertext.length, decryptedText, 0);

        return new String(decryptedText, StandardCharsets.UTF_8);
    }

    public static byte[] generateRandomKey() {
        byte[] key = new byte[32];
        SecureRandom random = new SecureRandom();
        random.nextBytes(key);
        return key;
    }

    public static byte[] generateRandomNonce() {
        byte[] nonce = new byte[8];
        SecureRandom random = new SecureRandom();
        random.nextBytes(nonce);
        return nonce;
    }

    public static String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(String.format("%02x", b & 0xff));
        }
        return sb.toString();
    }
}